"""
该文件仅用于对轨迹进行非极大值抑制
Use to NMST
"""


# 获取轨迹的所有子集切片
def subset(length, stride, batch_size):
    index = []
    for i in range(length - stride):
        for j in range(i + stride, length + 1):
            index.append([i, j])
    # 补齐BATCH_SIZE的整数倍长度，轨迹长度是动态的，但是BATCH_SIZE是静态的
    for i in range(batch_size - len(index) % batch_size):
        index.append([0, 0])
    return index


# 对轨迹聚合，已经把论文中的算法简化了
def merge(cluster):
    merged = []
    i = 0
    while i < len(cluster):
        j = 0
        while j < len(merged):
            # cluster: 0-预测值；1-预测概率；2-内圈起点；3-内圈终点；4-距离
            # merged: 0-预测值；1-内圈起点；2-内圈终点；3-外圈起点；4-外圈终点；5-累积概率；6-最大概率；7-累计距离；8-最大概率对应距离
            if cluster[i][0] == merged[j][0] and (cluster[i][2] == merged[j][1] or cluster[i][2] == merged[j][1] + 1)\
                    and merged[j][3] - 1 <= cluster[i][3] <= merged[j][4] + 1:
                merged[j][5] += cluster[i][1]
                if merged[j][7] < cluster[i][4]:
                    merged[j][7] = cluster[i][4]
                if cluster[i][1] > merged[j][6] or (cluster[i][1] == merged[j][6] and cluster[i][4] > merged[j][7]):
                    merged[j][6] = cluster[i][1]
                    merged[j][8] = cluster[i][4]
                if cluster[i][2] < merged[j][3]:
                    merged[j][3] = cluster[i][2]
                elif cluster[i][2] > merged[j][3]:
                    merged[j][1] = cluster[i][2]
                if cluster[i][3] < merged[j][2]:
                    merged[j][2] = cluster[i][3]
                elif cluster[i][3] > merged[j][2]:
                    merged[j][4] = cluster[i][3]
                if merged[j][1] > merged[j][2]:
                    merged[j][1] = merged[j][2]
                break
            j += 1
        if j == len(merged):
            merged.append([cluster[i][0], cluster[i][2], cluster[i][3], cluster[i][2], cluster[i][3],
                           cluster[i][1], cluster[i][1], cluster[i][4], cluster[i][4]])
        i += 1
    return merged
